import json, sys
from jinja2 import Environment, FileSystemLoader
from collections import defaultdict
from selenium import webdriver
import time, hashlib, base64, requests

def html_set_data(texdata, outprefix):
    path_dir = '/' + '/'.join(outprefix.split('/')[1:-1])
    env = Environment(loader=FileSystemLoader('/home/yong_sun/bin/workflow/qcstat/'))
    template = env.get_template('tempQC.html')
    buf = template.render(texdata)
    html = f'{outprefix}.fastq.Stats.html'
    with open(html, "w", encoding="utf-8") as fp:
        fp.write(buf)
    chrome_options = webdriver.ChromeOptions()
    chrome_options.add_argument('--headless')
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--disable-dev-shm-usage')
    browser = webdriver.Chrome(options=chrome_options)
    browser.get(f'file://{html}')
    time.sleep(3)
    browser.quit()

    try:
        url = 'https://qyapi.weixin.qq.com/cgi-bin/webhook/send?key=c635edcf-1adb-4277-9152-060d661c20b7'
        name = outprefix.split('/')[-1]
        requests.post(url, json={"msgtype": "text", "text": {"content": f"{name}"}})
        with open('fastq.stats.png', 'rb') as fq:
            png = fq.read()
            md5 = hashlib.md5(png).hexdigest()
            b64 = base64.b64encode(png).decode('utf-8')
            data = {'msgtype': 'image','image': {'base64': b64,'md5': md5}}
            requests.post(url, json=data)
    except Exception as e:
        print(e)
        sys.exit('PNG is not exists!')

def generate1mm(seq):
    mismatch = {'A':'CTG', 'T':'ACG', 'C':'ATG', 'G':'ATC'}
    mm = set()
    [mm.add(seq[0:i]+j+seq[i+1:]) for i, e in enumerate(seq) for j in mismatch[e]]
    return mm

def jsonoutdata(jsonname, outprefix):
    outdata = {'tabledata1': [], 'tabledata2': []}
    with open(jsonname) as js, open(f'{outprefix}.fastq.Stats.xls', 'w') as wxls:
        pjs = json.load(js)

        sampleReads = defaultdict(int)
        sampleIndex = defaultdict(str)
        undetermined = defaultdict(int)
        totalReads = 0

        for lane in pjs['ConversionResults']:
            for sample in lane['DemuxResults']:
                sampleReads[sample['SampleId']] += sample['NumberReads']
                sampleIndex[sample['SampleId']] = sample['IndexMetrics'][0]['IndexSequence']
                totalReads += sample['NumberReads']
            undetermined['Undetermined'] += lane['Undetermined']['NumberReads']
        totalReads += undetermined['Undetermined']

        otherBarcode = defaultdict(int)
        bcs = set()
        for lane in pjs['UnknownBarcodes']:
            for k, v in lane['Barcodes'].items():
                otherBarcode[k] += v
                bcs.add(k)

        outdata['tabledata1'].append(
            { 
                'sampleID': 'sampleID',
                'tindex': 'index',
                'treadsNumber': 'readsNumber',
                'tmax1misIndex': 'max1misIndex',
                'tmax1misNumber': 'max1misNumber',
                #'tsum1misNumber': 'sum1misNumber'
            }
        )
        outdata['tabledata1'].append(
            {
                'sampleID': 'TotalReads',
                'tindex': 'TotalReads',
                'treadsNumber': f'{totalReads:,}',
                'tmax1misIndex': '',
                'tmax1misNumber': '',
                #'tsum1misNumber': ''
            }
        )
        outdata['tabledata1'].append(
            {
                'sampleID': 'Undetermined',
                'tindex': 'Undetermined',
                'treadsNumber': f'{undetermined["Undetermined"]:,}',
                'tmax1misIndex': '',
                'tmax1misNumber': '',
                #'tsum1misNumber': ''
            }
        )
        

        wxls.write('sampleID\tindex\treadsNumber\t\tmax1misIndex\tmax1misNumber\tsum1misNumber\n')
        for k, v in sampleReads.items():
            idx = sampleIndex[k]
            idxSets = generate1mm(idx)
            combined = idxSets & bcs
            maxnum = 0
            mmidx = ''
            if v > 15000000 or k.startswith('NC') or k.find('YF') > -1 or k.find('KY') > -1 or k.find('SY') > -1:
                bit15M = False
            else:
                bit15M = True
            if combined:
                summm = 0
                for i in combined:
                    if maxnum < otherBarcode[i]:
                        maxnum = otherBarcode[i]
                        mmidx = i
                    summm += otherBarcode[i]
                wxls.write(f'{k}\t{idx}\t{v:,}\t\t{mmidx}\t{maxnum:,}\t{summm:,}\n')
                outdata['tabledata1'].append(
                  {
                    'sampleID': k,
                    'tindex': idx,
                    'treadsNumber': f'{v:,}',
                    'tmax1misIndex': mmidx,
                    'tmax1misNumber': f'{maxnum:,}',
                    #'tsum1misNumber': f'{summm:,}',
                    'bit15M': bit15M
                  }
                )
            else:
                wxls.write(f'{k}\t{idx}\t{v:,}\n')
                outdata['tabledata1'].append(
                  {
                    'sampleID': k,
                    'tindex': idx,
                    'treadsNumber': f'{v:,}',
                    'tmax1misIndex': mmidx,
                    'tmax1misNumber': '',
                    #'tsum1misNumber': '',
                    'bit15M': bit15M
                  }
                )
        
        n = 0
        wxls.write('\nUnknownBarcodes\treadsNumber\n')
        outdata['tabledata2'].append({ 'UnknownBarcodes': 'UnknownBarcodes','readsNumber': 'readsNumber' })
        for k, v in sorted(otherBarcode.items(), key=lambda x: x[1], reverse=True):
            n += 1
            if n < 10:
                wxls.write(f'{k}\t{v:,}\n')
                outdata['tabledata2'].append({ 'UnknownBarcodes': k,'readsNumber': f'{v:,}' })

    return outdata

if __name__ == '__main__':
    dataobj = jsonoutdata(sys.argv[1], sys.argv[2])
    html_set_data(dataobj, sys.argv[2])
